import cv2
import numpy as np

TRIANGLE = "triangle"
SQUARE   = "square"
CIRCLE   = "circle"

def extract_red(hsv):
    # 红色和蓝色的HSV阈值范围
    red_lower0 = np.array([0, 50, 50])
    red_upper0 = np.array([10, 255, 255])

    red_lower1 = np.array([156, 60, 60])
    red_upper1 = np.array([180, 255, 255])

    red_mask = cv2.inRange(hsv, red_lower0, red_upper0) + cv2.inRange(hsv, red_lower1, red_upper1)

    return red_mask


def extract_blue(hsv):
    blue_lower = np.array([100, 60, 60])
    blue_upper = np.array([124, 255, 255])

    # 从图像中分离红色和蓝色

    blue_mask = cv2.inRange(hsv, blue_lower, blue_upper)

    return blue_mask

def detect_shapes(bin):    
    epsilon = 0.05  # 多边形逼近程度
    '''
    检测二值图中的三角形，矩形，和圆形
    返回值list of (shape, area, (c0, c1))
    '''
    res = []
    # 查找轮廓
    contours, _ = cv2.findContours(bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 遍历所有轮廓
    for contour in contours:
        # 近似检测到的形状
        approx = cv2.approxPolyDP(contour, epsilon * cv2.arcLength(contour, True), True)

        # 忽略过小或非凸形状
        area = cv2.contourArea(contour)
        if area < 100 or not cv2.isContourConvex(approx):
            continue

        #获取坐标值和宽度、高度
        x, y, w, h = cv2.boundingRect(approx) 

        # 确定形状类型
        corners = len(approx)
        type = ""
        if corners == 3:
            type = TRIANGLE
        elif corners == 4:
            type = SQUARE
        else:
            type = CIRCLE
        
        shape = (type, area, (x+(w//2), y+(h//2)))
        res.append(shape)

        # debug
        cv2.drawContours(img, contour, -1, (0, 0, 255), 4)  #绘制轮廓线
        cv2.rectangle(img, (x,y),(x+w,y+h), (0, 255, 0), 2)  #绘制边界框
        cv2.putText(img, type, (x+(w//2), y+(h//2)), cv2.FONT_HERSHEY_COMPLEX, 0.6, (0,0,0), 1)  #绘制文字

    return res

# 读取图像
img = cv2.imread('./res/sample.jpg')

# 将图像转换为HSV颜色空间
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

r = extract_red(hsv)

shapes = detect_shapes(r)
print(shapes)

cv2.imshow("img", img)
cv2.imshow("color", r)

cv2.waitKey()
cv2.destroyAllWindows()
